{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# 2D 卷积优化\n",
        "\n",
        "**原作者**: [Thierry Moreau](https://homes.cs.washington.edu/~moreau/)\n",
        "\n",
        "本教程提供了关于如何使用 TVM 映射二维卷积工作负载有效的 VTA 设计的概述。建议先学习 {ref}`vta-mat-mult-opt` 教程。\n",
        "\n",
        "二维卷积在大多数计算机视觉深度神经网络中占主导地位。在本教程中，将演示 TVM 调度优化，将 NCHW 布局中的 2D 卷积算子映射到 VTA。还引入了延迟隐藏（latency hiding）的概念，它允许最大化 VTA 的计算和内存资源利用。\n",
        "\n",
        "## RPC 设置\n",
        "\n",
        "首先编程 Pynq 的 FPGA 并构建它的 RPC 运行时。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "tags": [
          "remove-cell"
        ]
      },
      "outputs": [],
      "source": [
        "import set_env"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import os\n",
        "import tvm\n",
        "import tvm.testing\n",
        "from tvm import te\n",
        "import vta\n",
        "import numpy as np\n",
        "\n",
        "from tvm import rpc\n",
        "from tvm.contrib import utils\n",
        "from vta.testing import simulator\n",
        "\n",
        "# Load VTA parameters from the 3rdparty/vta-hw/config/vta_config.json file\n",
        "env = vta.get_env()\n",
        "\n",
        "# We read the Pynq RPC host IP address and port number from the OS environment\n",
        "host = os.environ.get(\"VTA_RPC_HOST\", \"192.168.2.99\")\n",
        "port = int(os.environ.get(\"VTA_RPC_PORT\", \"9091\"))\n",
        "\n",
        "# We configure both the bitstream and the runtime system on the Pynq\n",
        "# to match the VTA configuration specified by the vta_config.json file.\n",
        "if env.TARGET == \"pynq\":\n",
        "    # Make sure that TVM was compiled with RPC=1\n",
        "    assert tvm.runtime.enabled(\"rpc\")\n",
        "    remote = rpc.connect(host, port)\n",
        "\n",
        "    # Reconfigure the JIT runtime\n",
        "    vta.reconfig_runtime(remote)\n",
        "\n",
        "    # Program the FPGA with a pre-compiled VTA bitstream.\n",
        "    # You can program the FPGA with your own custom bitstream\n",
        "    # by passing the path to the bitstream file instead of None.\n",
        "    vta.program_fpga(remote, bitstream=None)\n",
        "\n",
        "# In simulation mode, host the RPC server locally.\n",
        "elif env.TARGET in [\"sim\", \"tsim\"]:\n",
        "    remote = rpc.LocalSession()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 声明计算\n",
        "\n",
        "作为第一步，需要用 NCHW 格式描述 2D 卷积计算。\n",
        "\n",
        "通过 batch size、空间维度、输入通道、输出通道、核维度、核维度、填充维度和步长维度来定义二维卷积形状。\n",
        "\n",
        "选择 ResNet-18 架构的第 9 个卷积层的形状作为卷积 workload 参数。\n",
        "\n",
        "在 2D 卷积中添加了额外的算子，用于对输出进行移位和剪切，以模拟定点卷积之后的修正线性激活。将二维卷积层的 TVM 数据流图描述如下：\n",
        "\n",
        "```{image} images/conv2d_dataflow.png\n",
        ":align: center\n",
        "```\n",
        "\n",
        "这个计算被故意设置得太大，以至于不能一次全部放入 VTA 的 on-chip buffers。因此，在调度阶段，将依靠计算分块策略将计算分解为可管理的块。\n",
        "\n",
        "````{admonition} 空间填充\n",
        ":class: alert alert-info\n",
        "\n",
        "注意，需要导入 TOPI 库来对输入特征映射张量应用空间填充（Spatial padding）。空间填充有助于在 2D 卷积环境中分块，因为如果卷积核窗口大小大于 1，那么任何给定层的输入特征映射的相同 `(x, y)` 空间位置将被读取多次。在 CPU 和 GPU 上，当并行工作时，提高内存访问效率的一种方法是空间打包（spatial packing），这需要重新布局数据。VTA load DMA 引擎可以自动插入填充，这样原始的输入特征映射就不必在内存中重新打包。\n",
        "\n",
        "当数据从 DRAM load 到 VTA 的 SRAM 时，下面展示了 VTA 对动态空间填充的影响，随后是 2D 跨步和填充内存读取。\n",
        "\n",
        "```{image} images/padding.png\n",
        ":align: center\n",
        ":width: 480px\n",
        "```\n",
        "````"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from tvm import topi\n",
        "\n",
        "# 2D convolution layer dimensions taken from ResNet-18 architecture\n",
        "# (9th convolutional layer)\n",
        "batch_size = 1\n",
        "height = 14\n",
        "width = 14\n",
        "in_channels = 256\n",
        "out_channels = 256\n",
        "kernel_h = 3\n",
        "kernel_w = 3\n",
        "pad_h = 1\n",
        "pad_w = 1\n",
        "stride_h = 1\n",
        "stride_w = 1\n",
        "assert batch_size % env.BATCH == 0\n",
        "assert in_channels % env.BLOCK_IN == 0\n",
        "assert out_channels % env.BLOCK_OUT == 0\n",
        "\n",
        "# Input feature map: (N, IC, H, W, n, ic)\n",
        "data_shape = (\n",
        "    batch_size // env.BATCH,\n",
        "    in_channels // env.BLOCK_IN,\n",
        "    height,\n",
        "    width,\n",
        "    env.BATCH,\n",
        "    env.BLOCK_IN,\n",
        ")\n",
        "# Kernel: (OC, IC, H, W, oc, ic)\n",
        "kernel_shape = (\n",
        "    out_channels // env.BLOCK_OUT,\n",
        "    in_channels // env.BLOCK_IN,\n",
        "    kernel_h,\n",
        "    kernel_w,\n",
        "    env.BLOCK_OUT,\n",
        "    env.BLOCK_IN,\n",
        ")\n",
        "# Derive output feature map dimensions\n",
        "fout_height = (height + 2 * pad_h - kernel_h) // stride_h + 1\n",
        "fout_width = (width + 2 * pad_w - kernel_w) // stride_w + 1\n",
        "# Output feature map: (N, OC, H, W, n, oc)\n",
        "output_shape = (\n",
        "    batch_size // env.BATCH,\n",
        "    out_channels // env.BLOCK_OUT,\n",
        "    fout_height,\n",
        "    fout_width,\n",
        "    env.BATCH,\n",
        "    env.BLOCK_OUT,\n",
        ")\n",
        "\n",
        "# Convolution reduction axes\n",
        "dy = te.reduce_axis((0, kernel_h), name=\"dy\")\n",
        "dx = te.reduce_axis((0, kernel_w), name=\"dx\")\n",
        "ic = te.reduce_axis((0, in_channels // env.BLOCK_IN), name=\"ic\")\n",
        "ic_tns = te.reduce_axis((0, env.BLOCK_IN), name=\"ic_tns\")\n",
        "\n",
        "# Input placeholder tensors\n",
        "data = te.placeholder(data_shape, name=\"data\", dtype=env.inp_dtype)\n",
        "kernel = te.placeholder(kernel_shape, name=\"kernel\", dtype=env.wgt_dtype)\n",
        "\n",
        "# Copy buffers:\n",
        "#   Apply spatial padding to input feature map\n",
        "data_buf = topi.nn.pad(data, [0, 0, pad_h, pad_w, 0, 0], name=\"data_buf\")\n",
        "kernel_buf = te.compute(kernel_shape, lambda *i: kernel(*i), \"kernel_buf\")\n",
        "\n",
        "# Declare 2D convolution\n",
        "res_conv = te.compute(\n",
        "    output_shape,\n",
        "    lambda bo, co, i, j, bi, ci: te.sum(\n",
        "        data_buf[bo, ic, i * stride_h + dy, j * stride_w + dx, bi, ic_tns].astype(env.acc_dtype)\n",
        "        * kernel_buf[co, ic, dy, dx, ci, ic_tns].astype(env.acc_dtype),\n",
        "        axis=[ic, dy, dx, ic_tns],\n",
        "    ),\n",
        "    name=\"res_conv\",\n",
        ")\n",
        "\n",
        "# Add shift stage for fix-point normalization\n",
        "res_shr = te.compute(output_shape, lambda *i: res_conv(*i) >> 8, name=\"res_shr\")\n",
        "\n",
        "# Apply clipping between (0, input max value)\n",
        "inp_max = (1 << (env.INP_WIDTH - 1)) - 1\n",
        "res_max = te.compute(output_shape, lambda *i: tvm.te.max(res_shr(*i), 0), \"res_max\")\n",
        "res_min = te.compute(output_shape, lambda *i: tvm.te.min(res_max(*i), inp_max), \"res_min\")\n",
        "\n",
        "# Result Tensor\n",
        "res = te.compute(output_shape, lambda *i: res_min(*i).astype(env.inp_dtype), name=\"res\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 调度计算\n",
        "\n",
        "将看到一组必要的调度变换，以有效的方式将 2D 卷积映射到 VTA。这些包括：\n",
        "\n",
        "- 分块计算\n",
        "- 增加计算利用率（compute utilization）的虚拟线程（Virtual threading）\n",
        "- Lowering 到 VTA 硬件 intrinsics"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "collapsed": false
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
              "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
              "\n",
              "<span style=\"color: #A2F\">@I</span><span style=\"color: #A2F; font-weight: bold\">.</span>ir_module\n",
              "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #00F; font-weight: bold\">Module</span>:\n",
              "    <span style=\"color: #A2F\">@T</span><span style=\"color: #A2F; font-weight: bold\">.</span>prim_func\n",
              "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">main</span>(data: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">14</span>, <span style=\"color: #008000\">14</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>), <span style=\"color: #BA2121\">&quot;int8&quot;</span>), kernel: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">16</span>), <span style=\"color: #BA2121\">&quot;int8&quot;</span>), res: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">14</span>, <span style=\"color: #008000\">14</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>), <span style=\"color: #BA2121\">&quot;int8&quot;</span>)):\n",
              "        T<span style=\"color: #A2F; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;from_legacy_te_schedule&quot;</span>: T<span style=\"color: #A2F; font-weight: bold\">.</span>bool(<span style=\"color: #008000; font-weight: bold\">True</span>), <span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: T<span style=\"color: #A2F; font-weight: bold\">.</span>bool(<span style=\"color: #008000; font-weight: bold\">True</span>)})\n",
              "        data_buf <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>allocate([<span style=\"color: #008000\">65536</span>], <span style=\"color: #BA2121\">&quot;int8&quot;</span>, <span style=\"color: #BA2121\">&quot;global&quot;</span>)\n",
              "        kernel_buf <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>allocate([<span style=\"color: #008000\">589824</span>], <span style=\"color: #BA2121\">&quot;int8&quot;</span>, <span style=\"color: #BA2121\">&quot;global&quot;</span>)\n",
              "        res_conv <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>allocate([<span style=\"color: #008000\">50176</span>], <span style=\"color: #BA2121\">&quot;int32&quot;</span>, <span style=\"color: #BA2121\">&quot;global&quot;</span>)\n",
              "        data_buf_1 <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">65536</span>,), <span style=\"color: #BA2121\">&quot;int8&quot;</span>, data<span style=\"color: #A2F; font-weight: bold\">=</span>data_buf)\n",
              "        <span style=\"color: #008000; font-weight: bold\">for</span> i1, i2, i3, i5 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>grid(<span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">16</span>):\n",
              "            cse_var_1: T<span style=\"color: #A2F; font-weight: bold\">.</span>int32 <span style=\"color: #A2F; font-weight: bold\">=</span> i3 <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">16</span>\n",
              "            data_1 <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">50176</span>,), <span style=\"color: #BA2121\">&quot;int8&quot;</span>, data<span style=\"color: #A2F; font-weight: bold\">=</span>data<span style=\"color: #A2F; font-weight: bold\">.</span>data)\n",
              "            data_buf_1[i1 <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">4096</span> <span style=\"color: #A2F; font-weight: bold\">+</span> i2 <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">256</span> <span style=\"color: #A2F; font-weight: bold\">+</span> cse_var_1 <span style=\"color: #A2F; font-weight: bold\">+</span> i5] <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>if_then_else(<span style=\"color: #008000\">1</span> <span style=\"color: #A2F; font-weight: bold\">&lt;=</span> i2 <span style=\"color: #008000; font-weight: bold\">and</span> i2 <span style=\"color: #A2F; font-weight: bold\">&lt;</span> <span style=\"color: #008000\">15</span> <span style=\"color: #008000; font-weight: bold\">and</span> <span style=\"color: #008000\">1</span> <span style=\"color: #A2F; font-weight: bold\">&lt;=</span> i3 <span style=\"color: #008000; font-weight: bold\">and</span> i3 <span style=\"color: #A2F; font-weight: bold\">&lt;</span> <span style=\"color: #008000\">15</span>, data_1[i1 <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">3136</span> <span style=\"color: #A2F; font-weight: bold\">+</span> i2 <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">224</span> <span style=\"color: #A2F; font-weight: bold\">+</span> cse_var_1 <span style=\"color: #A2F; font-weight: bold\">+</span> i5 <span style=\"color: #A2F; font-weight: bold\">-</span> <span style=\"color: #008000\">240</span>], T<span style=\"color: #A2F; font-weight: bold\">.</span>int8(<span style=\"color: #008000\">0</span>))\n",
              "        kernel_buf_1 <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">589824</span>,), <span style=\"color: #BA2121\">&quot;int8&quot;</span>, data<span style=\"color: #A2F; font-weight: bold\">=</span>kernel_buf)\n",
              "        <span style=\"color: #008000; font-weight: bold\">for</span> i0, i1, i2, i3, i4, i5 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>grid(<span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">16</span>):\n",
              "            cse_var_2: T<span style=\"color: #A2F; font-weight: bold\">.</span>int32 <span style=\"color: #A2F; font-weight: bold\">=</span> i0 <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">36864</span> <span style=\"color: #A2F; font-weight: bold\">+</span> i1 <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">2304</span> <span style=\"color: #A2F; font-weight: bold\">+</span> i2 <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">768</span> <span style=\"color: #A2F; font-weight: bold\">+</span> i3 <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">256</span> <span style=\"color: #A2F; font-weight: bold\">+</span> i4 <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">16</span> <span style=\"color: #A2F; font-weight: bold\">+</span> i5\n",
              "            kernel_1 <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">589824</span>,), <span style=\"color: #BA2121\">&quot;int8&quot;</span>, data<span style=\"color: #A2F; font-weight: bold\">=</span>kernel<span style=\"color: #A2F; font-weight: bold\">.</span>data)\n",
              "            kernel_buf_1[cse_var_2] <span style=\"color: #A2F; font-weight: bold\">=</span> kernel_1[cse_var_2]\n",
              "        res_conv_1 <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">50176</span>,), <span style=\"color: #BA2121\">&quot;int32&quot;</span>, data<span style=\"color: #A2F; font-weight: bold\">=</span>res_conv)\n",
              "        <span style=\"color: #008000; font-weight: bold\">for</span> co, i, j, ci <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>grid(<span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">14</span>, <span style=\"color: #008000\">14</span>, <span style=\"color: #008000\">16</span>):\n",
              "            res_conv_1[co <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">3136</span> <span style=\"color: #A2F; font-weight: bold\">+</span> i <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">224</span> <span style=\"color: #A2F; font-weight: bold\">+</span> j <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">16</span> <span style=\"color: #A2F; font-weight: bold\">+</span> ci] <span style=\"color: #A2F; font-weight: bold\">=</span> <span style=\"color: #008000\">0</span>\n",
              "            <span style=\"color: #008000; font-weight: bold\">for</span> ic, dy, dx, ic_tns <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>grid(<span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">16</span>):\n",
              "                cse_var_4: T<span style=\"color: #A2F; font-weight: bold\">.</span>int32 <span style=\"color: #A2F; font-weight: bold\">=</span> j <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">16</span>\n",
              "                cse_var_3: T<span style=\"color: #A2F; font-weight: bold\">.</span>int32 <span style=\"color: #A2F; font-weight: bold\">=</span> co <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">3136</span> <span style=\"color: #A2F; font-weight: bold\">+</span> i <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">224</span> <span style=\"color: #A2F; font-weight: bold\">+</span> cse_var_4 <span style=\"color: #A2F; font-weight: bold\">+</span> ci\n",
              "                res_conv_1[cse_var_3] <span style=\"color: #A2F; font-weight: bold\">=</span> res_conv_1[cse_var_3] <span style=\"color: #A2F; font-weight: bold\">+</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>Cast(<span style=\"color: #BA2121\">&quot;int32&quot;</span>, data_buf_1[ic <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">4096</span> <span style=\"color: #A2F; font-weight: bold\">+</span> i <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">256</span> <span style=\"color: #A2F; font-weight: bold\">+</span> dy <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">256</span> <span style=\"color: #A2F; font-weight: bold\">+</span> cse_var_4 <span style=\"color: #A2F; font-weight: bold\">+</span> dx <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">16</span> <span style=\"color: #A2F; font-weight: bold\">+</span> ic_tns]) <span style=\"color: #A2F; font-weight: bold\">*</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>Cast(<span style=\"color: #BA2121\">&quot;int32&quot;</span>, kernel_buf_1[co <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">36864</span> <span style=\"color: #A2F; font-weight: bold\">+</span> ic <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">2304</span> <span style=\"color: #A2F; font-weight: bold\">+</span> dy <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">768</span> <span style=\"color: #A2F; font-weight: bold\">+</span> dx <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">256</span> <span style=\"color: #A2F; font-weight: bold\">+</span> ci <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">16</span> <span style=\"color: #A2F; font-weight: bold\">+</span> ic_tns])\n",
              "        res_conv_2 <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">50176</span>,), <span style=\"color: #BA2121\">&quot;int32&quot;</span>, data<span style=\"color: #A2F; font-weight: bold\">=</span>res_conv)\n",
              "        <span style=\"color: #008000; font-weight: bold\">for</span> i1, i2, i3, i5 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>grid(<span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">14</span>, <span style=\"color: #008000\">14</span>, <span style=\"color: #008000\">16</span>):\n",
              "            cse_var_5: T<span style=\"color: #A2F; font-weight: bold\">.</span>int32 <span style=\"color: #A2F; font-weight: bold\">=</span> i1 <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">3136</span> <span style=\"color: #A2F; font-weight: bold\">+</span> i2 <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">224</span> <span style=\"color: #A2F; font-weight: bold\">+</span> i3 <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">16</span> <span style=\"color: #A2F; font-weight: bold\">+</span> i5\n",
              "            res_conv_2[cse_var_5] <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>shift_right(res_conv_1[cse_var_5], <span style=\"color: #008000\">8</span>)\n",
              "        res_conv_3 <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">50176</span>,), <span style=\"color: #BA2121\">&quot;int32&quot;</span>, data<span style=\"color: #A2F; font-weight: bold\">=</span>res_conv)\n",
              "        <span style=\"color: #008000; font-weight: bold\">for</span> i1, i2, i3, i5 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>grid(<span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">14</span>, <span style=\"color: #008000\">14</span>, <span style=\"color: #008000\">16</span>):\n",
              "            cse_var_6: T<span style=\"color: #A2F; font-weight: bold\">.</span>int32 <span style=\"color: #A2F; font-weight: bold\">=</span> i1 <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">3136</span> <span style=\"color: #A2F; font-weight: bold\">+</span> i2 <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">224</span> <span style=\"color: #A2F; font-weight: bold\">+</span> i3 <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">16</span> <span style=\"color: #A2F; font-weight: bold\">+</span> i5\n",
              "            res_conv_3[cse_var_6] <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>max(res_conv_2[cse_var_6], <span style=\"color: #008000\">0</span>)\n",
              "        res_conv_4 <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">50176</span>,), <span style=\"color: #BA2121\">&quot;int32&quot;</span>, data<span style=\"color: #A2F; font-weight: bold\">=</span>res_conv)\n",
              "        <span style=\"color: #008000; font-weight: bold\">for</span> i1, i2, i3, i5 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>grid(<span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">14</span>, <span style=\"color: #008000\">14</span>, <span style=\"color: #008000\">16</span>):\n",
              "            cse_var_7: T<span style=\"color: #A2F; font-weight: bold\">.</span>int32 <span style=\"color: #A2F; font-weight: bold\">=</span> i1 <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">3136</span> <span style=\"color: #A2F; font-weight: bold\">+</span> i2 <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">224</span> <span style=\"color: #A2F; font-weight: bold\">+</span> i3 <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">16</span> <span style=\"color: #A2F; font-weight: bold\">+</span> i5\n",
              "            res_conv_4[cse_var_7] <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>min(res_conv_3[cse_var_7], <span style=\"color: #008000\">127</span>)\n",
              "        <span style=\"color: #008000; font-weight: bold\">for</span> i1, i2, i3, i5 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>grid(<span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">14</span>, <span style=\"color: #008000\">14</span>, <span style=\"color: #008000\">16</span>):\n",
              "            cse_var_8: T<span style=\"color: #A2F; font-weight: bold\">.</span>int32 <span style=\"color: #A2F; font-weight: bold\">=</span> i1 <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">3136</span> <span style=\"color: #A2F; font-weight: bold\">+</span> i2 <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">224</span> <span style=\"color: #A2F; font-weight: bold\">+</span> i3 <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">16</span> <span style=\"color: #A2F; font-weight: bold\">+</span> i5\n",
              "            res_1 <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">50176</span>,), <span style=\"color: #BA2121\">&quot;int8&quot;</span>, data<span style=\"color: #A2F; font-weight: bold\">=</span>res<span style=\"color: #A2F; font-weight: bold\">.</span>data)\n",
              "            res_1[cse_var_8] <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>Cast(<span style=\"color: #BA2121\">&quot;int8&quot;</span>, res_conv_4[cse_var_8])\n",
              "</pre></div>\n"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "# Create TVM schedule\n",
        "s = te.create_schedule(res.op)\n",
        "# Let's look at the default TVM schedule\n",
        "tvm.lower(s, [data, kernel, res], simple_mode=True).show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### 分块计算\n",
        "\n",
        "默认情况下，2D 卷积太大，激活或卷积核权重无法同时适应 VTA 的 on-chip buffer。沿着输入通道、输出通道和高度空间维度应用分块。不沿宽度空间维度进行分块，因为它是 NCHW 布局中的最内层维度（因此，为了增加局部性，最好不要沿最内层维度进行分块）。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {},
      "outputs": [],
      "source": [
        "# 定义 tiling sizes\n",
        "b_block = 1 // env.BATCH\n",
        "oc_block = 128 // env.BLOCK_OUT\n",
        "ic_block = 16 // env.BLOCK_IN\n",
        "h_block = 7\n",
        "w_block = 14\n",
        "\n",
        "# 沿着空间和输出通道维度平铺输出张量（因为默认情况下做单批推理，沿着批维分割没有效果）\n",
        "b, oc, y, x, b_tns, oc_tns = s[res].op.axis\n",
        "b_out, b_inn = s[res].split(b, factor=b_block)\n",
        "oc_out, oc_inn = s[res].split(oc, factor=oc_block)\n",
        "y_out, y_inn = s[res].split(y, factor=h_block)\n",
        "x_out, x_inn = s[res].split(x, factor=w_block)\n",
        "s[res].reorder(b_out, oc_out, y_out, x_out, b_inn, oc_inn, y_inn, x_inn, b_tns, oc_tns)\n",
        "\n",
        "# 将中间计算移动到每个输出计算 tile 中\n",
        "s[res_conv].compute_at(s[res], x_out)\n",
        "s[res_shr].compute_at(s[res], x_out)\n",
        "s[res_max].compute_at(s[res], x_out)\n",
        "s[res_min].compute_at(s[res], x_out)\n",
        "\n",
        "# 沿着规约轴（输入通道）应用额外的循环分割（loop split）\n",
        "b_inn, oc_inn, y_inn, x_inn, b_tns, oc_tns = s[res_conv].op.axis\n",
        "ic_out, ic_inn = s[res_conv].split(ic, factor=ic_block)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### 重排轴\n",
        "\n",
        "1. 将 VTA 张量轴分组在最内侧位置：b_tns, oc_tns, ic_tns，允许 TVM 张量化。\n",
        "2. 将 ic_out 轴整个移出卷积循环，沿着归约轴进行分块。\n",
        "3. 重新排列块轴：b_inn, oc_inn, y_inn, x_inn, ic_inn, dy, dx。VTA 运行时/硬件要求为每个 VTA 张量运算写入不同的输出特征映射（feature map）位置。这个限制要求我们将 oc_inn、y_inn 或 x_inn 中的一个排序在 b_tns 之前，因为它们都会影响输出特征映射索引。因此，我们选择把 x_inn 放在里面，如下所示。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "s[res_conv].reorder(ic_out, b_inn, oc_inn, y_inn, ic_inn, dy, dx, x_inn, b_tns, oc_tns, ic_tns)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### 虚拟线程\n",
        "\n",
        "虚拟线程（virtual thread）是一种在 VTA 硬件设计中增加任务级管道并行性的机制。换句话说，它通过隐藏内存访问延迟（hiding memory access latency）提高了计算资源的利用率。\n",
        "\n",
        "在下面的实现中，虚拟线程将工作分配给沿输出通道轴划分的两个线程。在下面的图中，展示了计算 2D 卷积时工作是如何分割的。\n",
        "\n",
        "```{image} images/virtual_threading.png\n",
        ":align: center\n",
        ":width: 480px\n",
        "```"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "collapsed": false
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
              "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
              "\n",
              "<span style=\"color: #A2F\">@I</span><span style=\"color: #A2F; font-weight: bold\">.</span>ir_module\n",
              "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #00F; font-weight: bold\">Module</span>:\n",
              "    <span style=\"color: #A2F\">@T</span><span style=\"color: #A2F; font-weight: bold\">.</span>prim_func\n",
              "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">main</span>(data: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">14</span>, <span style=\"color: #008000\">14</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>), <span style=\"color: #BA2121\">&quot;int8&quot;</span>), kernel: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">16</span>), <span style=\"color: #BA2121\">&quot;int8&quot;</span>), res: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">14</span>, <span style=\"color: #008000\">14</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>), <span style=\"color: #BA2121\">&quot;int8&quot;</span>)):\n",
              "        T<span style=\"color: #A2F; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;from_legacy_te_schedule&quot;</span>: T<span style=\"color: #A2F; font-weight: bold\">.</span>bool(<span style=\"color: #008000; font-weight: bold\">True</span>), <span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: T<span style=\"color: #A2F; font-weight: bold\">.</span>bool(<span style=\"color: #008000; font-weight: bold\">True</span>)})\n",
              "        data_buf <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>allocate([<span style=\"color: #008000\">65536</span>], <span style=\"color: #BA2121\">&quot;int8&quot;</span>, <span style=\"color: #BA2121\">&quot;global&quot;</span>)\n",
              "        kernel_buf <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>allocate([<span style=\"color: #008000\">589824</span>], <span style=\"color: #BA2121\">&quot;int8&quot;</span>, <span style=\"color: #BA2121\">&quot;global&quot;</span>)\n",
              "        res_conv <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>allocate([<span style=\"color: #008000\">25088</span>], <span style=\"color: #BA2121\">&quot;int32&quot;</span>, <span style=\"color: #BA2121\">&quot;global&quot;</span>)\n",
              "        data_buf_1 <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">65536</span>,), <span style=\"color: #BA2121\">&quot;int8&quot;</span>, data<span style=\"color: #A2F; font-weight: bold\">=</span>data_buf)\n",
              "        <span style=\"color: #008000; font-weight: bold\">for</span> i1, i2, i3, i5 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>grid(<span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">16</span>):\n",
              "            cse_var_1: T<span style=\"color: #A2F; font-weight: bold\">.</span>int32 <span style=\"color: #A2F; font-weight: bold\">=</span> i3 <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">16</span>\n",
              "            data_1 <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">50176</span>,), <span style=\"color: #BA2121\">&quot;int8&quot;</span>, data<span style=\"color: #A2F; font-weight: bold\">=</span>data<span style=\"color: #A2F; font-weight: bold\">.</span>data)\n",
              "            data_buf_1[i1 <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">4096</span> <span style=\"color: #A2F; font-weight: bold\">+</span> i2 <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">256</span> <span style=\"color: #A2F; font-weight: bold\">+</span> cse_var_1 <span style=\"color: #A2F; font-weight: bold\">+</span> i5] <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>if_then_else(<span style=\"color: #008000\">1</span> <span style=\"color: #A2F; font-weight: bold\">&lt;=</span> i2 <span style=\"color: #008000; font-weight: bold\">and</span> i2 <span style=\"color: #A2F; font-weight: bold\">&lt;</span> <span style=\"color: #008000\">15</span> <span style=\"color: #008000; font-weight: bold\">and</span> <span style=\"color: #008000\">1</span> <span style=\"color: #A2F; font-weight: bold\">&lt;=</span> i3 <span style=\"color: #008000; font-weight: bold\">and</span> i3 <span style=\"color: #A2F; font-weight: bold\">&lt;</span> <span style=\"color: #008000\">15</span>, data_1[i1 <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">3136</span> <span style=\"color: #A2F; font-weight: bold\">+</span> i2 <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">224</span> <span style=\"color: #A2F; font-weight: bold\">+</span> cse_var_1 <span style=\"color: #A2F; font-weight: bold\">+</span> i5 <span style=\"color: #A2F; font-weight: bold\">-</span> <span style=\"color: #008000\">240</span>], T<span style=\"color: #A2F; font-weight: bold\">.</span>int8(<span style=\"color: #008000\">0</span>))\n",
              "        kernel_buf_1 <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">589824</span>,), <span style=\"color: #BA2121\">&quot;int8&quot;</span>, data<span style=\"color: #A2F; font-weight: bold\">=</span>kernel_buf)\n",
              "        <span style=\"color: #008000; font-weight: bold\">for</span> i0, i1, i2, i3, i4, i5 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>grid(<span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">16</span>):\n",
              "            cse_var_2: T<span style=\"color: #A2F; font-weight: bold\">.</span>int32 <span style=\"color: #A2F; font-weight: bold\">=</span> i0 <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">36864</span> <span style=\"color: #A2F; font-weight: bold\">+</span> i1 <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">2304</span> <span style=\"color: #A2F; font-weight: bold\">+</span> i2 <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">768</span> <span style=\"color: #A2F; font-weight: bold\">+</span> i3 <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">256</span> <span style=\"color: #A2F; font-weight: bold\">+</span> i4 <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">16</span> <span style=\"color: #A2F; font-weight: bold\">+</span> i5\n",
              "            kernel_1 <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">589824</span>,), <span style=\"color: #BA2121\">&quot;int8&quot;</span>, data<span style=\"color: #A2F; font-weight: bold\">=</span>kernel<span style=\"color: #A2F; font-weight: bold\">.</span>data)\n",
              "            kernel_buf_1[cse_var_2] <span style=\"color: #A2F; font-weight: bold\">=</span> kernel_1[cse_var_2]\n",
              "        <span style=\"color: #008000; font-weight: bold\">for</span> i2_outer <span style=\"color: #008000; font-weight: bold\">in</span> range(<span style=\"color: #008000\">2</span>):\n",
              "            res_conv_1 <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">157351936</span>,), <span style=\"color: #BA2121\">&quot;int32&quot;</span>, data<span style=\"color: #A2F; font-weight: bold\">=</span>res_conv)\n",
              "            <span style=\"color: #008000; font-weight: bold\">for</span> co_init, i_init, j_init, ci_init <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>grid(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">7</span>, <span style=\"color: #008000\">14</span>, <span style=\"color: #008000\">16</span>):\n",
              "                cse_var_3: T<span style=\"color: #A2F; font-weight: bold\">.</span>int32 <span style=\"color: #A2F; font-weight: bold\">=</span> co_init <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">1568</span> <span style=\"color: #A2F; font-weight: bold\">+</span> i_init <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">224</span> <span style=\"color: #A2F; font-weight: bold\">+</span> j_init <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">16</span> <span style=\"color: #A2F; font-weight: bold\">+</span> ci_init\n",
              "                res_conv_1[cse_var_3] <span style=\"color: #A2F; font-weight: bold\">=</span> <span style=\"color: #008000\">0</span>\n",
              "                res_conv_1[cse_var_3 <span style=\"color: #A2F; font-weight: bold\">+</span> <span style=\"color: #008000\">12544</span>] <span style=\"color: #A2F; font-weight: bold\">=</span> <span style=\"color: #008000\">0</span>\n",
              "            <span style=\"color: #008000; font-weight: bold\">for</span> ic_outer, co, i, dy, dx, j, ci, ic_tns <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>grid(<span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">7</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">14</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">16</span>):\n",
              "                cse_var_8: T<span style=\"color: #A2F; font-weight: bold\">.</span>int32 <span style=\"color: #A2F; font-weight: bold\">=</span> j <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">16</span>\n",
              "                cse_var_7: T<span style=\"color: #A2F; font-weight: bold\">.</span>int32 <span style=\"color: #A2F; font-weight: bold\">=</span> co <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">1568</span> <span style=\"color: #A2F; font-weight: bold\">+</span> i <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">224</span> <span style=\"color: #A2F; font-weight: bold\">+</span> cse_var_8 <span style=\"color: #A2F; font-weight: bold\">+</span> ci\n",
              "                cse_var_6: T<span style=\"color: #A2F; font-weight: bold\">.</span>int32 <span style=\"color: #A2F; font-weight: bold\">=</span> cse_var_7 <span style=\"color: #A2F; font-weight: bold\">+</span> <span style=\"color: #008000\">12544</span>\n",
              "                cse_var_5: T<span style=\"color: #A2F; font-weight: bold\">.</span>int32 <span style=\"color: #A2F; font-weight: bold\">=</span> co <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">36864</span> <span style=\"color: #A2F; font-weight: bold\">+</span> ic_outer <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">2304</span> <span style=\"color: #A2F; font-weight: bold\">+</span> dy <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">768</span> <span style=\"color: #A2F; font-weight: bold\">+</span> dx <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">256</span> <span style=\"color: #A2F; font-weight: bold\">+</span> ci <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">16</span> <span style=\"color: #A2F; font-weight: bold\">+</span> ic_tns\n",
              "                cse_var_4: T<span style=\"color: #A2F; font-weight: bold\">.</span>int32 <span style=\"color: #A2F; font-weight: bold\">=</span> ic_outer <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">4096</span> <span style=\"color: #A2F; font-weight: bold\">+</span> i2_outer <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">1792</span> <span style=\"color: #A2F; font-weight: bold\">+</span> i <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">256</span> <span style=\"color: #A2F; font-weight: bold\">+</span> dy <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">256</span> <span style=\"color: #A2F; font-weight: bold\">+</span> cse_var_8 <span style=\"color: #A2F; font-weight: bold\">+</span> dx <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">16</span> <span style=\"color: #A2F; font-weight: bold\">+</span> ic_tns\n",
              "                res_conv_1[cse_var_7] <span style=\"color: #A2F; font-weight: bold\">=</span> res_conv_1[cse_var_7] <span style=\"color: #A2F; font-weight: bold\">+</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>Cast(<span style=\"color: #BA2121\">&quot;int32&quot;</span>, data_buf_1[cse_var_4]) <span style=\"color: #A2F; font-weight: bold\">*</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>Cast(<span style=\"color: #BA2121\">&quot;int32&quot;</span>, kernel_buf_1[cse_var_5])\n",
              "                res_conv_1[cse_var_6] <span style=\"color: #A2F; font-weight: bold\">=</span> res_conv_1[cse_var_6] <span style=\"color: #A2F; font-weight: bold\">+</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>Cast(<span style=\"color: #BA2121\">&quot;int32&quot;</span>, data_buf_1[cse_var_4]) <span style=\"color: #A2F; font-weight: bold\">*</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>Cast(<span style=\"color: #BA2121\">&quot;int32&quot;</span>, kernel_buf_1[cse_var_5 <span style=\"color: #A2F; font-weight: bold\">+</span> <span style=\"color: #008000\">294912</span>])\n",
              "            res_conv_2 <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">157351936</span>,), <span style=\"color: #BA2121\">&quot;int32&quot;</span>, data<span style=\"color: #A2F; font-weight: bold\">=</span>res_conv)\n",
              "            <span style=\"color: #008000; font-weight: bold\">for</span> i1, i2, i3, i5 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>grid(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">7</span>, <span style=\"color: #008000\">14</span>, <span style=\"color: #008000\">16</span>):\n",
              "                cse_var_10: T<span style=\"color: #A2F; font-weight: bold\">.</span>int32 <span style=\"color: #A2F; font-weight: bold\">=</span> i1 <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">1568</span> <span style=\"color: #A2F; font-weight: bold\">+</span> i2 <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">224</span> <span style=\"color: #A2F; font-weight: bold\">+</span> i3 <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">16</span> <span style=\"color: #A2F; font-weight: bold\">+</span> i5\n",
              "                cse_var_9: T<span style=\"color: #A2F; font-weight: bold\">.</span>int32 <span style=\"color: #A2F; font-weight: bold\">=</span> cse_var_10 <span style=\"color: #A2F; font-weight: bold\">+</span> <span style=\"color: #008000\">12544</span>\n",
              "                res_conv_2[cse_var_10] <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>shift_right(res_conv_1[cse_var_10], <span style=\"color: #008000\">8</span>)\n",
              "                res_conv_2[cse_var_9] <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>shift_right(res_conv_1[cse_var_9], <span style=\"color: #008000\">8</span>)\n",
              "            res_conv_3 <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">157351936</span>,), <span style=\"color: #BA2121\">&quot;int32&quot;</span>, data<span style=\"color: #A2F; font-weight: bold\">=</span>res_conv)\n",
              "            <span style=\"color: #008000; font-weight: bold\">for</span> i1, i2, i3, i5 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>grid(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">7</span>, <span style=\"color: #008000\">14</span>, <span style=\"color: #008000\">16</span>):\n",
              "                cse_var_12: T<span style=\"color: #A2F; font-weight: bold\">.</span>int32 <span style=\"color: #A2F; font-weight: bold\">=</span> i1 <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">1568</span> <span style=\"color: #A2F; font-weight: bold\">+</span> i2 <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">224</span> <span style=\"color: #A2F; font-weight: bold\">+</span> i3 <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">16</span> <span style=\"color: #A2F; font-weight: bold\">+</span> i5\n",
              "                cse_var_11: T<span style=\"color: #A2F; font-weight: bold\">.</span>int32 <span style=\"color: #A2F; font-weight: bold\">=</span> cse_var_12 <span style=\"color: #A2F; font-weight: bold\">+</span> <span style=\"color: #008000\">12544</span>\n",
              "                res_conv_3[cse_var_12] <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>max(res_conv_2[cse_var_12], <span style=\"color: #008000\">0</span>)\n",
              "                res_conv_3[cse_var_11] <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>max(res_conv_2[cse_var_11], <span style=\"color: #008000\">0</span>)\n",
              "            res_conv_4 <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">157351936</span>,), <span style=\"color: #BA2121\">&quot;int32&quot;</span>, data<span style=\"color: #A2F; font-weight: bold\">=</span>res_conv)\n",
              "            <span style=\"color: #008000; font-weight: bold\">for</span> i1, i2, i3, i5 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>grid(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">7</span>, <span style=\"color: #008000\">14</span>, <span style=\"color: #008000\">16</span>):\n",
              "                cse_var_14: T<span style=\"color: #A2F; font-weight: bold\">.</span>int32 <span style=\"color: #A2F; font-weight: bold\">=</span> i1 <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">1568</span> <span style=\"color: #A2F; font-weight: bold\">+</span> i2 <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">224</span> <span style=\"color: #A2F; font-weight: bold\">+</span> i3 <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">16</span> <span style=\"color: #A2F; font-weight: bold\">+</span> i5\n",
              "                cse_var_13: T<span style=\"color: #A2F; font-weight: bold\">.</span>int32 <span style=\"color: #A2F; font-weight: bold\">=</span> cse_var_14 <span style=\"color: #A2F; font-weight: bold\">+</span> <span style=\"color: #008000\">12544</span>\n",
              "                res_conv_4[cse_var_14] <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>min(res_conv_3[cse_var_14], <span style=\"color: #008000\">127</span>)\n",
              "                res_conv_4[cse_var_13] <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>min(res_conv_3[cse_var_13], <span style=\"color: #008000\">127</span>)\n",
              "            <span style=\"color: #008000; font-weight: bold\">for</span> i1_inner, i2_inner, i3_inner, i5 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>grid(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">7</span>, <span style=\"color: #008000\">14</span>, <span style=\"color: #008000\">16</span>):\n",
              "                cse_var_18: T<span style=\"color: #A2F; font-weight: bold\">.</span>int32 <span style=\"color: #A2F; font-weight: bold\">=</span> i2_inner <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">224</span>\n",
              "                cse_var_17: T<span style=\"color: #A2F; font-weight: bold\">.</span>int32 <span style=\"color: #A2F; font-weight: bold\">=</span> i3_inner <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">16</span>\n",
              "                cse_var_16: T<span style=\"color: #A2F; font-weight: bold\">.</span>int32 <span style=\"color: #A2F; font-weight: bold\">=</span> i1_inner <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">1568</span> <span style=\"color: #A2F; font-weight: bold\">+</span> cse_var_18 <span style=\"color: #A2F; font-weight: bold\">+</span> cse_var_17 <span style=\"color: #A2F; font-weight: bold\">+</span> i5\n",
              "                cse_var_15: T<span style=\"color: #A2F; font-weight: bold\">.</span>int32 <span style=\"color: #A2F; font-weight: bold\">=</span> i1_inner <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">3136</span> <span style=\"color: #A2F; font-weight: bold\">+</span> i2_outer <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">1568</span> <span style=\"color: #A2F; font-weight: bold\">+</span> cse_var_18 <span style=\"color: #A2F; font-weight: bold\">+</span> cse_var_17 <span style=\"color: #A2F; font-weight: bold\">+</span> i5\n",
              "                res_1 <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">50176</span>,), <span style=\"color: #BA2121\">&quot;int8&quot;</span>, data<span style=\"color: #A2F; font-weight: bold\">=</span>res<span style=\"color: #A2F; font-weight: bold\">.</span>data)\n",
              "                res_1[cse_var_15] <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>Cast(<span style=\"color: #BA2121\">&quot;int8&quot;</span>, res_conv_4[cse_var_16])\n",
              "                res_1[cse_var_15 <span style=\"color: #A2F; font-weight: bold\">+</span> <span style=\"color: #008000\">25088</span>] <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>Cast(<span style=\"color: #BA2121\">&quot;int8&quot;</span>, res_conv_4[cse_var_16 <span style=\"color: #A2F; font-weight: bold\">+</span> <span style=\"color: #008000\">12544</span>])\n",
              "</pre></div>\n"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "# VTA 仅支持 2 个虚拟线程\n",
        "v_threads = 2\n",
        "\n",
        "# 沿输出通道外轴执行虚拟线程 split \n",
        "_, tx = s[res].split(oc_out, factor=v_threads)\n",
        "s[res].reorder(tx, b_out)\n",
        "s[res].bind(tx, te.thread_axis(\"cthread\"))\n",
        "# 看看在分块和虚拟线程之后当前的 TVM 调度\n",
        "tvm.lower(s, [data, kernel, res], simple_mode=True).show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Lowering Copies 到 DMA Transfers\n",
        "\n",
        "接下来，设置相应的 on-chip VTA SRAM buffers 的 buffers 作用域。将 load 循环移动到 2D 卷积计算循环，以 stage 内存加载，以便它们适合 on-chip SRAM buffers。最后，用 DMA 复制 pragma 注解了 load/store 循环外轴，以便在 VTA 上执行大容量内存传输。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Set scope of SRAM buffers\n",
        "s[data_buf].set_scope(env.inp_scope)\n",
        "s[kernel_buf].set_scope(env.wgt_scope)\n",
        "s[res_conv].set_scope(env.acc_scope)\n",
        "s[res_shr].set_scope(env.acc_scope)\n",
        "s[res_min].set_scope(env.acc_scope)\n",
        "s[res_max].set_scope(env.acc_scope)\n",
        "\n",
        "# Block data and kernel cache reads\n",
        "s[data_buf].compute_at(s[res_conv], ic_out)\n",
        "s[kernel_buf].compute_at(s[res_conv], ic_out)\n",
        "\n",
        "# Use DMA copy pragma on DRAM->SRAM operations\n",
        "s[data_buf].pragma(s[data_buf].op.axis[0], env.dma_copy)\n",
        "s[kernel_buf].pragma(s[kernel_buf].op.axis[0], env.dma_copy)\n",
        "\n",
        "# 在每个结果块中对 SRAM->DRAM 操作使用 DMA copy pragma（这意味着这些 copy 应该沿着 b_inn 或结果轴 4 执行）\n",
        "s[res].pragma(s[res].op.axis[4], env.dma_copy)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Lowering 计算到 VTA 计算 Intrinsics\n",
        "\n",
        "最后阶段是通过将二维卷积映射为张量 intrinsics，并将位移和剪切计算映射为向量 ALU，从而将计算循环 lower 到 VTA 硬件 intrinsics。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Apply tensorization over the batch tensor tile axis\n",
        "s[res_conv].tensorize(b_tns, env.gemm)\n",
        "\n",
        "# Add an ALU pragma over the shift and clipping operations\n",
        "s[res_shr].pragma(s[res_shr].op.axis[0], env.alu)\n",
        "s[res_min].pragma(s[res_min].op.axis[0], env.alu)\n",
        "s[res_max].pragma(s[res_max].op.axis[0], env.alu)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "看看在将内存 loads/stores lower 到 DMA copy intrinsic，并将计算 lowering 到 VTA 计算 intrinsic 之后，最终 lowering 的 TVM 调度。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "collapsed": false
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "[20:15:05] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/tir/transforms/arg_binder.cc:95: Warning: Trying to bind buffer to another one with lower alignment requirement  required_alignment=256, provided_alignment=64\n",
            "[20:15:05] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/script/printer/tir/expr.cc:246: Warning: No TScriptPrinterName attribute for tir.vta.coproc_dep_push\n",
            "[20:15:05] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/script/printer/tir/expr.cc:246: Warning: No TScriptPrinterName attribute for tir.vta.coproc_dep_push\n",
            "[20:15:05] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/script/printer/tir/expr.cc:246: Warning: No TScriptPrinterName attribute for tir.vta.coproc_dep_pop\n",
            "[20:15:05] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/script/printer/tir/expr.cc:246: Warning: No TScriptPrinterName attribute for tir.vta.uop_push\n",
            "[20:15:05] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/script/printer/tir/expr.cc:246: Warning: No TScriptPrinterName attribute for tir.vta.coproc_dep_push\n",
            "[20:15:05] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/script/printer/tir/expr.cc:246: Warning: No TScriptPrinterName attribute for tir.vta.coproc_dep_pop\n",
            "[20:15:05] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/script/printer/tir/expr.cc:246: Warning: No TScriptPrinterName attribute for tir.vta.command_handle\n",
            "[20:15:05] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/script/printer/tir/expr.cc:246: Warning: No TScriptPrinterName attribute for tir.vta.command_handle\n",
            "[20:15:05] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/script/printer/tir/expr.cc:246: Warning: No TScriptPrinterName attribute for tir.vta.coproc_dep_push\n",
            "[20:15:05] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/script/printer/tir/expr.cc:246: Warning: No TScriptPrinterName attribute for tir.vta.coproc_dep_pop\n",
            "[20:15:05] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/script/printer/tir/expr.cc:246: Warning: No TScriptPrinterName attribute for tir.vta.command_handle\n",
            "[20:15:05] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/script/printer/tir/expr.cc:246: Warning: No TScriptPrinterName attribute for tir.vta.command_handle\n",
            "[20:15:05] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/script/printer/tir/expr.cc:246: Warning: No TScriptPrinterName attribute for tir.vta.coproc_dep_push\n",
            "[20:15:05] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/script/printer/tir/expr.cc:246: Warning: No TScriptPrinterName attribute for tir.vta.coproc_dep_pop\n",
            "[20:15:05] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/script/printer/tir/expr.cc:246: Warning: No TScriptPrinterName attribute for tir.vta.uop_push\n",
            "[20:15:05] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/script/printer/tir/expr.cc:246: Warning: No TScriptPrinterName attribute for tir.vta.coproc_dep_push\n",
            "[20:15:05] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/script/printer/tir/expr.cc:246: Warning: No TScriptPrinterName attribute for tir.vta.coproc_dep_pop\n",
            "[20:15:05] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/script/printer/tir/expr.cc:246: Warning: No TScriptPrinterName attribute for tir.vta.coproc_dep_pop\n",
            "[20:15:05] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/script/printer/tir/expr.cc:246: Warning: No TScriptPrinterName attribute for tir.vta.uop_push\n",
            "[20:15:05] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/script/printer/tir/expr.cc:246: Warning: No TScriptPrinterName attribute for tir.vta.uop_push\n",
            "[20:15:05] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/script/printer/tir/expr.cc:246: Warning: No TScriptPrinterName attribute for tir.vta.uop_push\n",
            "[20:15:05] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/script/printer/tir/expr.cc:246: Warning: No TScriptPrinterName attribute for tir.vta.coproc_dep_push\n",
            "[20:15:05] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/script/printer/tir/expr.cc:246: Warning: No TScriptPrinterName attribute for tir.vta.coproc_dep_pop\n",
            "[20:15:05] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/script/printer/tir/expr.cc:246: Warning: No TScriptPrinterName attribute for tir.vta.command_handle\n",
            "[20:15:05] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/script/printer/tir/expr.cc:246: Warning: No TScriptPrinterName attribute for tir.vta.coproc_dep_push\n",
            "[20:15:05] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/script/printer/tir/expr.cc:246: Warning: No TScriptPrinterName attribute for tir.vta.coproc_dep_pop\n",
            "[20:15:05] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/script/printer/tir/expr.cc:246: Warning: No TScriptPrinterName attribute for tir.vta.coproc_dep_pop\n",
            "[20:15:05] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/script/printer/tir/expr.cc:246: Warning: No TScriptPrinterName attribute for tir.vta.coproc_sync\n"
          ]
        },
        {
          "data": {
            "text/html": [
              "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
              "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
              "\n",
              "<span style=\"color: #A2F\">@I</span><span style=\"color: #A2F; font-weight: bold\">.</span>ir_module\n",
              "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #00F; font-weight: bold\">Module</span>:\n",
              "    <span style=\"color: #A2F\">@T</span><span style=\"color: #A2F; font-weight: bold\">.</span>prim_func\n",
              "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">main</span>(data: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">14</span>, <span style=\"color: #008000\">14</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>), <span style=\"color: #BA2121\">&quot;int8&quot;</span>), kernel: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">16</span>), <span style=\"color: #BA2121\">&quot;int8&quot;</span>), res: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">14</span>, <span style=\"color: #008000\">14</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>), <span style=\"color: #BA2121\">&quot;int8&quot;</span>)):\n",
              "        T<span style=\"color: #A2F; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;from_legacy_te_schedule&quot;</span>: T<span style=\"color: #A2F; font-weight: bold\">.</span>bool(<span style=\"color: #008000; font-weight: bold\">True</span>), <span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: T<span style=\"color: #A2F; font-weight: bold\">.</span>bool(<span style=\"color: #008000; font-weight: bold\">True</span>)})\n",
              "        T<span style=\"color: #A2F; font-weight: bold\">.</span>tir<span style=\"color: #A2F; font-weight: bold\">.</span>vta<span style=\"color: #A2F; font-weight: bold\">.</span>coproc_dep_push(<span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">2</span>)\n",
              "        T<span style=\"color: #A2F; font-weight: bold\">.</span>tir<span style=\"color: #A2F; font-weight: bold\">.</span>vta<span style=\"color: #A2F; font-weight: bold\">.</span>coproc_dep_push(<span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">2</span>)\n",
              "        <span style=\"color: #008000; font-weight: bold\">for</span> i2_outer <span style=\"color: #008000; font-weight: bold\">in</span> range(<span style=\"color: #008000\">2</span>):\n",
              "            <span style=\"color: #008000; font-weight: bold\">for</span> cthread_s <span style=\"color: #008000; font-weight: bold\">in</span> range(<span style=\"color: #008000\">2</span>):\n",
              "                vta <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>int32()\n",
              "                T<span style=\"color: #A2F; font-weight: bold\">.</span>attr(T<span style=\"color: #A2F; font-weight: bold\">.</span>iter_var(vta, <span style=\"color: #008000; font-weight: bold\">None</span>, <span style=\"color: #BA2121\">&quot;ThreadIndex&quot;</span>, <span style=\"color: #BA2121\">&quot;vta&quot;</span>), <span style=\"color: #BA2121\">&quot;coproc_scope&quot;</span>, <span style=\"color: #008000\">2</span>)\n",
              "                T<span style=\"color: #A2F; font-weight: bold\">.</span>tir<span style=\"color: #A2F; font-weight: bold\">.</span>vta<span style=\"color: #A2F; font-weight: bold\">.</span>coproc_dep_pop(<span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">2</span>)\n",
              "                <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>attr(T<span style=\"color: #A2F; font-weight: bold\">.</span>iter_var(vta, <span style=\"color: #008000; font-weight: bold\">None</span>, <span style=\"color: #BA2121\">&quot;ThreadIndex&quot;</span>, <span style=\"color: #BA2121\">&quot;vta&quot;</span>), <span style=\"color: #BA2121\">&quot;coproc_uop_scope&quot;</span>, <span style=\"color: #BA2121\">&quot;VTAPushGEMMOp&quot;</span>):\n",
              "                    T<span style=\"color: #A2F; font-weight: bold\">.</span>call_extern(<span style=\"color: #BA2121\">&quot;int32&quot;</span>, <span style=\"color: #BA2121\">&quot;VTAUopLoopBegin&quot;</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">98</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>)\n",
              "                    T<span style=\"color: #A2F; font-weight: bold\">.</span>call_extern(<span style=\"color: #BA2121\">&quot;int32&quot;</span>, <span style=\"color: #BA2121\">&quot;VTAUopLoopBegin&quot;</span>, <span style=\"color: #008000\">7</span>, <span style=\"color: #008000\">14</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>)\n",
              "                    <span style=\"color: #008000; font-weight: bold\">for</span> j_init <span style=\"color: #008000; font-weight: bold\">in</span> range(<span style=\"color: #008000\">14</span>):\n",
              "                        T<span style=\"color: #A2F; font-weight: bold\">.</span>tir<span style=\"color: #A2F; font-weight: bold\">.</span>vta<span style=\"color: #A2F; font-weight: bold\">.</span>uop_push(<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">1</span>, cthread_s <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">784</span> <span style=\"color: #A2F; font-weight: bold\">+</span> j_init, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>)\n",
              "                    T<span style=\"color: #A2F; font-weight: bold\">.</span>call_extern(<span style=\"color: #BA2121\">&quot;int32&quot;</span>, <span style=\"color: #BA2121\">&quot;VTAUopLoopEnd&quot;</span>)\n",
              "                    T<span style=\"color: #A2F; font-weight: bold\">.</span>call_extern(<span style=\"color: #BA2121\">&quot;int32&quot;</span>, <span style=\"color: #BA2121\">&quot;VTAUopLoopEnd&quot;</span>)\n",
              "                T<span style=\"color: #A2F; font-weight: bold\">.</span>tir<span style=\"color: #A2F; font-weight: bold\">.</span>vta<span style=\"color: #A2F; font-weight: bold\">.</span>coproc_dep_push(<span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">1</span>)\n",
              "            vta <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>int32()\n",
              "            <span style=\"color: #008000; font-weight: bold\">for</span> ic_outer <span style=\"color: #008000; font-weight: bold\">in</span> range(<span style=\"color: #008000\">16</span>):\n",
              "                cse_var_6: T<span style=\"color: #A2F; font-weight: bold\">.</span>int32 <span style=\"color: #A2F; font-weight: bold\">=</span> i2_outer <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">7</span>\n",
              "                cse_var_5: T<span style=\"color: #A2F; font-weight: bold\">.</span>int32 <span style=\"color: #A2F; font-weight: bold\">=</span> ic_outer <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">9</span>\n",
              "                cse_var_4: T<span style=\"color: #A2F; font-weight: bold\">.</span>int32 <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>max(<span style=\"color: #008000\">1</span> <span style=\"color: #A2F; font-weight: bold\">-</span> cse_var_6, <span style=\"color: #008000\">0</span>)\n",
              "                cse_var_3: T<span style=\"color: #A2F; font-weight: bold\">.</span>int32 <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>max(cse_var_6 <span style=\"color: #A2F; font-weight: bold\">-</span> <span style=\"color: #008000\">6</span>, <span style=\"color: #008000\">0</span>)\n",
              "                cse_var_2: T<span style=\"color: #A2F; font-weight: bold\">.</span>int32 <span style=\"color: #A2F; font-weight: bold\">=</span> <span style=\"color: #008000\">9</span> <span style=\"color: #A2F; font-weight: bold\">-</span> cse_var_4 <span style=\"color: #A2F; font-weight: bold\">-</span> cse_var_3\n",
              "                cse_var_1: T<span style=\"color: #A2F; font-weight: bold\">.</span>int32 <span style=\"color: #A2F; font-weight: bold\">=</span> ic_outer <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">196</span> <span style=\"color: #A2F; font-weight: bold\">+</span> i2_outer <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">98</span> <span style=\"color: #A2F; font-weight: bold\">+</span> cse_var_4 <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">14</span> <span style=\"color: #A2F; font-weight: bold\">-</span> <span style=\"color: #008000\">14</span>\n",
              "                <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>attr(T<span style=\"color: #A2F; font-weight: bold\">.</span>iter_var(vta, <span style=\"color: #008000; font-weight: bold\">None</span>, <span style=\"color: #BA2121\">&quot;ThreadIndex&quot;</span>, <span style=\"color: #BA2121\">&quot;vta&quot;</span>), <span style=\"color: #BA2121\">&quot;coproc_scope&quot;</span>, <span style=\"color: #008000\">1</span>):\n",
              "                    T<span style=\"color: #A2F; font-weight: bold\">.</span>tir<span style=\"color: #A2F; font-weight: bold\">.</span>vta<span style=\"color: #A2F; font-weight: bold\">.</span>coproc_dep_pop(<span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">1</span>)\n",
              "                    T<span style=\"color: #A2F; font-weight: bold\">.</span>call_extern(<span style=\"color: #BA2121\">&quot;int32&quot;</span>, <span style=\"color: #BA2121\">&quot;VTALoadBuffer2D&quot;</span>, T<span style=\"color: #A2F; font-weight: bold\">.</span>tvm_thread_context(T<span style=\"color: #A2F; font-weight: bold\">.</span>tir<span style=\"color: #A2F; font-weight: bold\">.</span>vta<span style=\"color: #A2F; font-weight: bold\">.</span>command_handle()), data<span style=\"color: #A2F; font-weight: bold\">.</span>data, cse_var_1, <span style=\"color: #008000\">14</span>, cse_var_2, <span style=\"color: #008000\">14</span>, <span style=\"color: #008000\">1</span>, cse_var_4, <span style=\"color: #008000\">1</span>, cse_var_3, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">2</span>)\n",
              "                    T<span style=\"color: #A2F; font-weight: bold\">.</span>call_extern(<span style=\"color: #BA2121\">&quot;int32&quot;</span>, <span style=\"color: #BA2121\">&quot;VTALoadBuffer2D&quot;</span>, T<span style=\"color: #A2F; font-weight: bold\">.</span>tvm_thread_context(T<span style=\"color: #A2F; font-weight: bold\">.</span>tir<span style=\"color: #A2F; font-weight: bold\">.</span>vta<span style=\"color: #A2F; font-weight: bold\">.</span>command_handle()), kernel<span style=\"color: #A2F; font-weight: bold\">.</span>data, cse_var_5, <span style=\"color: #008000\">9</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">144</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">1</span>)\n",
              "                    T<span style=\"color: #A2F; font-weight: bold\">.</span>tir<span style=\"color: #A2F; font-weight: bold\">.</span>vta<span style=\"color: #A2F; font-weight: bold\">.</span>coproc_dep_push(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">2</span>)\n",
              "                <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>attr(T<span style=\"color: #A2F; font-weight: bold\">.</span>iter_var(vta, <span style=\"color: #008000; font-weight: bold\">None</span>, <span style=\"color: #BA2121\">&quot;ThreadIndex&quot;</span>, <span style=\"color: #BA2121\">&quot;vta&quot;</span>), <span style=\"color: #BA2121\">&quot;coproc_scope&quot;</span>, <span style=\"color: #008000\">1</span>):\n",
              "                    T<span style=\"color: #A2F; font-weight: bold\">.</span>tir<span style=\"color: #A2F; font-weight: bold\">.</span>vta<span style=\"color: #A2F; font-weight: bold\">.</span>coproc_dep_pop(<span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">1</span>)\n",
              "                    T<span style=\"color: #A2F; font-weight: bold\">.</span>call_extern(<span style=\"color: #BA2121\">&quot;int32&quot;</span>, <span style=\"color: #BA2121\">&quot;VTALoadBuffer2D&quot;</span>, T<span style=\"color: #A2F; font-weight: bold\">.</span>tvm_thread_context(T<span style=\"color: #A2F; font-weight: bold\">.</span>tir<span style=\"color: #A2F; font-weight: bold\">.</span>vta<span style=\"color: #A2F; font-weight: bold\">.</span>command_handle()), data<span style=\"color: #A2F; font-weight: bold\">.</span>data, cse_var_1, <span style=\"color: #008000\">14</span>, cse_var_2, <span style=\"color: #008000\">14</span>, <span style=\"color: #008000\">1</span>, cse_var_4, <span style=\"color: #008000\">1</span>, cse_var_3, <span style=\"color: #008000\">144</span>, <span style=\"color: #008000\">2</span>)\n",
              "                    T<span style=\"color: #A2F; font-weight: bold\">.</span>call_extern(<span style=\"color: #BA2121\">&quot;int32&quot;</span>, <span style=\"color: #BA2121\">&quot;VTALoadBuffer2D&quot;</span>, T<span style=\"color: #A2F; font-weight: bold\">.</span>tvm_thread_context(T<span style=\"color: #A2F; font-weight: bold\">.</span>tir<span style=\"color: #A2F; font-weight: bold\">.</span>vta<span style=\"color: #A2F; font-weight: bold\">.</span>command_handle()), kernel<span style=\"color: #A2F; font-weight: bold\">.</span>data, cse_var_5 <span style=\"color: #A2F; font-weight: bold\">+</span> <span style=\"color: #008000\">1152</span>, <span style=\"color: #008000\">9</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">144</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">72</span>, <span style=\"color: #008000\">1</span>)\n",
              "                    T<span style=\"color: #A2F; font-weight: bold\">.</span>tir<span style=\"color: #A2F; font-weight: bold\">.</span>vta<span style=\"color: #A2F; font-weight: bold\">.</span>coproc_dep_push(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">2</span>)\n",
              "                <span style=\"color: #008000; font-weight: bold\">for</span> cthread_s <span style=\"color: #008000; font-weight: bold\">in</span> range(<span style=\"color: #008000\">2</span>):\n",
              "                    T<span style=\"color: #A2F; font-weight: bold\">.</span>attr(T<span style=\"color: #A2F; font-weight: bold\">.</span>iter_var(vta, <span style=\"color: #008000; font-weight: bold\">None</span>, <span style=\"color: #BA2121\">&quot;ThreadIndex&quot;</span>, <span style=\"color: #BA2121\">&quot;vta&quot;</span>), <span style=\"color: #BA2121\">&quot;coproc_scope&quot;</span>, <span style=\"color: #008000\">2</span>)\n",
              "                    T<span style=\"color: #A2F; font-weight: bold\">.</span>tir<span style=\"color: #A2F; font-weight: bold\">.</span>vta<span style=\"color: #A2F; font-weight: bold\">.</span>coproc_dep_pop(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">2</span>)\n",
              "                    <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>attr(T<span style=\"color: #A2F; font-weight: bold\">.</span>iter_var(vta, <span style=\"color: #008000; font-weight: bold\">None</span>, <span style=\"color: #BA2121\">&quot;ThreadIndex&quot;</span>, <span style=\"color: #BA2121\">&quot;vta&quot;</span>), <span style=\"color: #BA2121\">&quot;coproc_uop_scope&quot;</span>, <span style=\"color: #BA2121\">&quot;VTAPushGEMMOp&quot;</span>):\n",
              "                        T<span style=\"color: #A2F; font-weight: bold\">.</span>call_extern(<span style=\"color: #BA2121\">&quot;int32&quot;</span>, <span style=\"color: #BA2121\">&quot;VTAUopLoopBegin&quot;</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">98</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">9</span>)\n",
              "                        T<span style=\"color: #A2F; font-weight: bold\">.</span>call_extern(<span style=\"color: #BA2121\">&quot;int32&quot;</span>, <span style=\"color: #BA2121\">&quot;VTAUopLoopBegin&quot;</span>, <span style=\"color: #008000\">7</span>, <span style=\"color: #008000\">14</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">0</span>)\n",
              "                        <span style=\"color: #008000; font-weight: bold\">for</span> dy, dx, j <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>grid(<span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">14</span>):\n",
              "                            T<span style=\"color: #A2F; font-weight: bold\">.</span>tir<span style=\"color: #A2F; font-weight: bold\">.</span>vta<span style=\"color: #A2F; font-weight: bold\">.</span>uop_push(<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, cthread_s <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">784</span> <span style=\"color: #A2F; font-weight: bold\">+</span> j, cthread_s <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">144</span> <span style=\"color: #A2F; font-weight: bold\">+</span> dy <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">16</span> <span style=\"color: #A2F; font-weight: bold\">+</span> j <span style=\"color: #A2F; font-weight: bold\">+</span> dx, cthread_s <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">72</span> <span style=\"color: #A2F; font-weight: bold\">+</span> dy <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">3</span> <span style=\"color: #A2F; font-weight: bold\">+</span> dx, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>)\n",
              "                        T<span style=\"color: #A2F; font-weight: bold\">.</span>call_extern(<span style=\"color: #BA2121\">&quot;int32&quot;</span>, <span style=\"color: #BA2121\">&quot;VTAUopLoopEnd&quot;</span>)\n",
              "                        T<span style=\"color: #A2F; font-weight: bold\">.</span>call_extern(<span style=\"color: #BA2121\">&quot;int32&quot;</span>, <span style=\"color: #BA2121\">&quot;VTAUopLoopEnd&quot;</span>)\n",
              "                    T<span style=\"color: #A2F; font-weight: bold\">.</span>tir<span style=\"color: #A2F; font-weight: bold\">.</span>vta<span style=\"color: #A2F; font-weight: bold\">.</span>coproc_dep_push(<span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">1</span>)\n",
              "            T<span style=\"color: #A2F; font-weight: bold\">.</span>tir<span style=\"color: #A2F; font-weight: bold\">.</span>vta<span style=\"color: #A2F; font-weight: bold\">.</span>coproc_dep_pop(<span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">1</span>)\n",
              "            T<span style=\"color: #A2F; font-weight: bold\">.</span>tir<span style=\"color: #A2F; font-weight: bold\">.</span>vta<span style=\"color: #A2F; font-weight: bold\">.</span>coproc_dep_pop(<span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">1</span>)\n",
              "            <span style=\"color: #008000; font-weight: bold\">for</span> cthread_s <span style=\"color: #008000; font-weight: bold\">in</span> range(<span style=\"color: #008000\">2</span>):\n",
              "                cse_var_7: T<span style=\"color: #A2F; font-weight: bold\">.</span>int32 <span style=\"color: #A2F; font-weight: bold\">=</span> cthread_s <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">784</span>\n",
              "                T<span style=\"color: #A2F; font-weight: bold\">.</span>attr(T<span style=\"color: #A2F; font-weight: bold\">.</span>iter_var(vta, <span style=\"color: #008000; font-weight: bold\">None</span>, <span style=\"color: #BA2121\">&quot;ThreadIndex&quot;</span>, <span style=\"color: #BA2121\">&quot;vta&quot;</span>), <span style=\"color: #BA2121\">&quot;coproc_scope&quot;</span>, <span style=\"color: #008000\">2</span>)\n",
              "                <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>attr(T<span style=\"color: #A2F; font-weight: bold\">.</span>iter_var(vta, <span style=\"color: #008000; font-weight: bold\">None</span>, <span style=\"color: #BA2121\">&quot;ThreadIndex&quot;</span>, <span style=\"color: #BA2121\">&quot;vta&quot;</span>), <span style=\"color: #BA2121\">&quot;coproc_uop_scope&quot;</span>, <span style=\"color: #BA2121\">&quot;VTAPushALUOp&quot;</span>):\n",
              "                    T<span style=\"color: #A2F; font-weight: bold\">.</span>call_extern(<span style=\"color: #BA2121\">&quot;int32&quot;</span>, <span style=\"color: #BA2121\">&quot;VTAUopLoopBegin&quot;</span>, <span style=\"color: #008000\">784</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">0</span>)\n",
              "                    T<span style=\"color: #A2F; font-weight: bold\">.</span>tir<span style=\"color: #A2F; font-weight: bold\">.</span>vta<span style=\"color: #A2F; font-weight: bold\">.</span>uop_push(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">0</span>, cse_var_7, cse_var_7, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">8</span>)\n",
              "                    T<span style=\"color: #A2F; font-weight: bold\">.</span>call_extern(<span style=\"color: #BA2121\">&quot;int32&quot;</span>, <span style=\"color: #BA2121\">&quot;VTAUopLoopEnd&quot;</span>)\n",
              "                <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>attr(T<span style=\"color: #A2F; font-weight: bold\">.</span>iter_var(vta, <span style=\"color: #008000; font-weight: bold\">None</span>, <span style=\"color: #BA2121\">&quot;ThreadIndex&quot;</span>, <span style=\"color: #BA2121\">&quot;vta&quot;</span>), <span style=\"color: #BA2121\">&quot;coproc_uop_scope&quot;</span>, <span style=\"color: #BA2121\">&quot;VTAPushALUOp&quot;</span>):\n",
              "                    T<span style=\"color: #A2F; font-weight: bold\">.</span>call_extern(<span style=\"color: #BA2121\">&quot;int32&quot;</span>, <span style=\"color: #BA2121\">&quot;VTAUopLoopBegin&quot;</span>, <span style=\"color: #008000\">784</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">0</span>)\n",
              "                    T<span style=\"color: #A2F; font-weight: bold\">.</span>tir<span style=\"color: #A2F; font-weight: bold\">.</span>vta<span style=\"color: #A2F; font-weight: bold\">.</span>uop_push(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">0</span>, cse_var_7, cse_var_7, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">0</span>)\n",
              "                    T<span style=\"color: #A2F; font-weight: bold\">.</span>call_extern(<span style=\"color: #BA2121\">&quot;int32&quot;</span>, <span style=\"color: #BA2121\">&quot;VTAUopLoopEnd&quot;</span>)\n",
              "                <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>attr(T<span style=\"color: #A2F; font-weight: bold\">.</span>iter_var(vta, <span style=\"color: #008000; font-weight: bold\">None</span>, <span style=\"color: #BA2121\">&quot;ThreadIndex&quot;</span>, <span style=\"color: #BA2121\">&quot;vta&quot;</span>), <span style=\"color: #BA2121\">&quot;coproc_uop_scope&quot;</span>, <span style=\"color: #BA2121\">&quot;VTAPushALUOp&quot;</span>):\n",
              "                    T<span style=\"color: #A2F; font-weight: bold\">.</span>call_extern(<span style=\"color: #BA2121\">&quot;int32&quot;</span>, <span style=\"color: #BA2121\">&quot;VTAUopLoopBegin&quot;</span>, <span style=\"color: #008000\">784</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">0</span>)\n",
              "                    T<span style=\"color: #A2F; font-weight: bold\">.</span>tir<span style=\"color: #A2F; font-weight: bold\">.</span>vta<span style=\"color: #A2F; font-weight: bold\">.</span>uop_push(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">0</span>, cse_var_7, cse_var_7, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">127</span>)\n",
              "                    T<span style=\"color: #A2F; font-weight: bold\">.</span>call_extern(<span style=\"color: #BA2121\">&quot;int32&quot;</span>, <span style=\"color: #BA2121\">&quot;VTAUopLoopEnd&quot;</span>)\n",
              "                T<span style=\"color: #A2F; font-weight: bold\">.</span>tir<span style=\"color: #A2F; font-weight: bold\">.</span>vta<span style=\"color: #A2F; font-weight: bold\">.</span>coproc_dep_push(<span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">3</span>)\n",
              "            <span style=\"color: #008000; font-weight: bold\">for</span> cthread_s <span style=\"color: #008000; font-weight: bold\">in</span> range(<span style=\"color: #008000\">2</span>):\n",
              "                T<span style=\"color: #A2F; font-weight: bold\">.</span>attr(T<span style=\"color: #A2F; font-weight: bold\">.</span>iter_var(vta, <span style=\"color: #008000; font-weight: bold\">None</span>, <span style=\"color: #BA2121\">&quot;ThreadIndex&quot;</span>, <span style=\"color: #BA2121\">&quot;vta&quot;</span>), <span style=\"color: #BA2121\">&quot;coproc_scope&quot;</span>, <span style=\"color: #008000\">3</span>)\n",
              "                T<span style=\"color: #A2F; font-weight: bold\">.</span>tir<span style=\"color: #A2F; font-weight: bold\">.</span>vta<span style=\"color: #A2F; font-weight: bold\">.</span>coproc_dep_pop(<span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">3</span>)\n",
              "                <span style=\"color: #008000; font-weight: bold\">for</span> i1_inner, i2_inner, i3_inner <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>grid(<span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">7</span>, <span style=\"color: #008000\">14</span>):\n",
              "                    cse_var_8: T<span style=\"color: #A2F; font-weight: bold\">.</span>int32 <span style=\"color: #A2F; font-weight: bold\">=</span> i2_inner <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">14</span>\n",
              "                    T<span style=\"color: #A2F; font-weight: bold\">.</span>call_extern(<span style=\"color: #BA2121\">&quot;int32&quot;</span>, <span style=\"color: #BA2121\">&quot;VTAStoreBuffer2D&quot;</span>, T<span style=\"color: #A2F; font-weight: bold\">.</span>tvm_thread_context(T<span style=\"color: #A2F; font-weight: bold\">.</span>tir<span style=\"color: #A2F; font-weight: bold\">.</span>vta<span style=\"color: #A2F; font-weight: bold\">.</span>command_handle()), cthread_s <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">784</span> <span style=\"color: #A2F; font-weight: bold\">+</span> i1_inner <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">98</span> <span style=\"color: #A2F; font-weight: bold\">+</span> cse_var_8 <span style=\"color: #A2F; font-weight: bold\">+</span> i3_inner, <span style=\"color: #008000\">4</span>, res<span style=\"color: #A2F; font-weight: bold\">.</span>data, cthread_s <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">1568</span> <span style=\"color: #A2F; font-weight: bold\">+</span> i1_inner <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">196</span> <span style=\"color: #A2F; font-weight: bold\">+</span> i2_outer <span style=\"color: #A2F; font-weight: bold\">*</span> <span style=\"color: #008000\">98</span> <span style=\"color: #A2F; font-weight: bold\">+</span> cse_var_8 <span style=\"color: #A2F; font-weight: bold\">+</span> i3_inner, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>)\n",
              "                T<span style=\"color: #A2F; font-weight: bold\">.</span>tir<span style=\"color: #A2F; font-weight: bold\">.</span>vta<span style=\"color: #A2F; font-weight: bold\">.</span>coproc_dep_push(<span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">2</span>)\n",
              "        T<span style=\"color: #A2F; font-weight: bold\">.</span>tir<span style=\"color: #A2F; font-weight: bold\">.</span>vta<span style=\"color: #A2F; font-weight: bold\">.</span>coproc_dep_pop(<span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">2</span>)\n",
              "        T<span style=\"color: #A2F; font-weight: bold\">.</span>tir<span style=\"color: #A2F; font-weight: bold\">.</span>vta<span style=\"color: #A2F; font-weight: bold\">.</span>coproc_dep_pop(<span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">2</span>)\n",
              "        T<span style=\"color: #A2F; font-weight: bold\">.</span>tir<span style=\"color: #A2F; font-weight: bold\">.</span>vta<span style=\"color: #A2F; font-weight: bold\">.</span>coproc_sync()\n",
              "</pre></div>\n"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "vta.lower(s, [data, kernel, res], simple_mode=True).show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## TVM 计算和验证\n",
        "\n",
        "在指定调度之后，可以将其编译为 TVM 函数。保存模块，这样就可以通过 RPC 发送它。运行该函数并对 numpy 实现进行验证，以确保其正确性。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "[20:15:06] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/tir/transforms/arg_binder.cc:95: Warning: Trying to bind buffer to another one with lower alignment requirement  required_alignment=256, provided_alignment=64\n",
            "2025-07-27 20:15:06.974 INFO load_module /tmp/tmpib410yd2/conv2d.o\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Execution statistics:\n",
            "\tinp_load_nbytes :           114688\n",
            "\twgt_load_nbytes :          1179648\n",
            "\tacc_load_nbytes :                0\n",
            "\tuop_load_nbytes :             1144\n",
            "\tout_store_nbytes:            50176\n",
            "\tgemm_counter    :           451584\n",
            "\talu_counter     :             9408\n",
            "Successful 2D convolution test!\n"
          ]
        }
      ],
      "source": [
        "# This library facilitates 2D convolution testing\n",
        "from tvm.topi.testing import conv2d_nchw_python\n",
        "\n",
        "# Compile the TVM module\n",
        "with vta.build_config(disabled_pass={\"tir.CommonSubexprElimTIR\"}):\n",
        "    my_conv = vta.build(\n",
        "        s, [data, kernel, res], tvm.target.Target(\"ext_dev\", host=env.target_host), name=\"my_conv\"\n",
        "    )\n",
        "temp = utils.tempdir()\n",
        "my_conv.save(temp.relpath(\"conv2d.o\"))\n",
        "remote.upload(temp.relpath(\"conv2d.o\"))\n",
        "f = remote.load_module(\"conv2d.o\")\n",
        "\n",
        "# Get the remote device context\n",
        "ctx = remote.ext_dev(0)\n",
        "\n",
        "# Initialize the data and kernel arrays randomly in the int range\n",
        "# of (-128, 128] in NCHW layout\n",
        "data_np = np.random.randint(-128, 128, size=(batch_size, in_channels, height, width)).astype(\n",
        "    data.dtype\n",
        ")\n",
        "kernel_np = np.random.randint(\n",
        "    -128, 128, size=(out_channels, in_channels, kernel_h, kernel_w)\n",
        ").astype(kernel.dtype)\n",
        "\n",
        "# Apply packing to the data and kernel arrays from a 2D NCHW\n",
        "# to a 4D NCHWnc packed layout\n",
        "data_packed = data_np.reshape(\n",
        "    batch_size // env.BATCH, env.BATCH, in_channels // env.BLOCK_IN, env.BLOCK_IN, height, width\n",
        ").transpose((0, 2, 4, 5, 1, 3))\n",
        "\n",
        "kernel_packed = kernel_np.reshape(\n",
        "    out_channels // env.BLOCK_OUT,\n",
        "    env.BLOCK_OUT,\n",
        "    in_channels // env.BLOCK_IN,\n",
        "    env.BLOCK_IN,\n",
        "    kernel_h,\n",
        "    kernel_w,\n",
        ").transpose((0, 2, 4, 5, 1, 3))\n",
        "\n",
        "# Format the input/output arrays with tvm.nd.array to the DLPack standard\n",
        "data_nd = tvm.nd.array(data_packed, ctx)\n",
        "kernel_nd = tvm.nd.array(kernel_packed, ctx)\n",
        "res_nd = tvm.nd.array(np.zeros(output_shape).astype(res.dtype), ctx)\n",
        "\n",
        "# Clear stats\n",
        "if env.TARGET in [\"sim\", \"tsim\"]:\n",
        "    simulator.clear_stats()\n",
        "\n",
        "# Invoke the module to perform the computation\n",
        "f(data_nd, kernel_nd, res_nd)\n",
        "\n",
        "# Verify against numpy implementation\n",
        "res_ref = conv2d_nchw_python(\n",
        "    data_np.astype(env.acc_dtype),\n",
        "    kernel_np.astype(env.acc_dtype),\n",
        "    (stride_h, stride_w),\n",
        "    (pad_h, pad_w),\n",
        ").astype(env.acc_dtype)\n",
        "res_ref = res_ref >> env.INP_WIDTH\n",
        "res_ref = np.clip(res_ref, 0, inp_max)\n",
        "res_ref = res_ref.astype(res.dtype)\n",
        "res_ref = res_ref.reshape(\n",
        "    (\n",
        "        batch_size // env.BATCH,\n",
        "        env.BATCH,\n",
        "        out_channels // env.BLOCK_OUT,\n",
        "        env.BLOCK_OUT,\n",
        "        fout_height,\n",
        "        fout_width,\n",
        "    )\n",
        ").transpose((0, 2, 4, 5, 1, 3))\n",
        "np.testing.assert_allclose(res_ref, res_nd.numpy())\n",
        "\n",
        "# Print stats\n",
        "if env.TARGET in [\"sim\", \"tsim\"]:\n",
        "    sim_stats = simulator.stats()\n",
        "    print(\"Execution statistics:\")\n",
        "    for k, v in sim_stats.items():\n",
        "        print(\"\\t{:<16}: {:>16}\".format(k, v))\n",
        "\n",
        "print(\"Successful 2D convolution test!\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "```{topic} 小结\n",
        "本教程演示如何使用 TVM 调度原语 lower 硬件加速器 intrinsics 的 2D 卷积，利用特定于硬件的优化，比如使用带虚拟线程的隐藏延迟。\n",
        "```"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.12.2"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
